package com.yang.sso.oauth2.oauth.customizer;

import org.yang.common.core.constant.JwtClaimConstants;
import com.yang.sso.oauth2.model.MyUserDetails;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.authority.AuthorityUtils;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;

import java.util.Collections;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * JWT 自定义字段配置
 *
 * @author: lslands
 * @description:
 * @version：v1.0
 * @date: 2024/6/3 12:16
 */
@Slf4j
@Configuration
public class JwtTokenCustomizerConfig {


    /**
     * JWT 自定义字段
     */
    @Bean
    public OAuth2TokenCustomizer<JwtEncodingContext> jwtTokenCustomizer() {
        return context -> {
            if (OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType()) && context.getPrincipal() instanceof UsernamePasswordAuthenticationToken) {
                // 自定义access_token的标题/声明
                Optional.ofNullable(context.getPrincipal().getPrincipal()).ifPresent(principal -> {
                    JwtClaimsSet.Builder claims = context.getClaims();
                    // 系统用户添加自定义字段
                    if (principal instanceof MyUserDetails userDetails) {
                        claims.claim(JwtClaimConstants.USER_ID, userDetails.getUserId());
                        claims.claim(JwtClaimConstants.USERNAME, userDetails.getUsername());
                        claims.claim(JwtClaimConstants.NICKNAME,userDetails.getNickname());
                        claims.claim(JwtClaimConstants.COMPANY_ID, userDetails.getCompanyId());
                        claims.claim(JwtClaimConstants.ENABLED,userDetails.getEnabled());
                        // 这里存入角色至JWT，解析JWT的角色用于鉴权的位置: ResourceServerConfig#jwtAuthenticationConverter
                        var authorities = AuthorityUtils.authorityListToSet(context.getPrincipal().getAuthorities())
                                .stream()
                                .collect(Collectors.collectingAndThen(Collectors.toSet(), Collections::unmodifiableSet));
                        claims.claim(JwtClaimConstants.AUTHORITIES, authorities);

                    }
                });
            }
        };
    }
}
